import tensorflow as tf
import numpy as np
import os
import math

# you need to change this to your data directory
train_dir = 'img/'


def get_files(pic_list, class_names):
    """
    Args:
        pic_list: include n_classes pics path
    Returns:
        list of images and labels
    """
    n_classes = len(pic_list)
    classes = []
    labels = []
    name_to_label = {}
    for i in range(n_classes):
        name_to_label[i] = class_names[i]
        classes.append(pic_list[i])
        label = []
        for j in range(len(pic_list[i])):
            label.append(i)
        labels.append(label)

    for i in range(n_classes):
        print(f'there are {len(labels[i])} {class_names[i]}s')

    image_list = []
    label_list = []
    for i in range(n_classes):
        image_list = np.hstack((image_list, classes[i]))
        label_list = np.hstack((label_list, labels[i]))

    temp = np.array([image_list, label_list])
    temp = temp.transpose()
    np.random.shuffle(temp)

    all_image_list = temp[:, 0]
    all_label_list = temp[:, 1]

    n_sample = len(all_label_list)
    n_val = math.ceil(n_sample * 0.2)  # number of validation samples
    n_train = n_sample - n_val  # number of trainning samples

    tra_images = all_image_list[0:n_train]
    tra_labels = all_label_list[0:n_train]
    tra_labels = [int(float(i)) for i in tra_labels]
    val_images = all_image_list[n_train:-1]
    val_labels = all_label_list[n_train:-1]
    val_labels = [int(float(i)) for i in val_labels]

    return tra_images, tra_labels, val_images, val_labels, name_to_label


def get_batch(image, label, image_W, image_H, batch_size, capacity):
    """
    Args:
        image: list type
        label: list type
        image_W: image width
        image_H: image height
        batch_size: batch size
        capacity: the maximum elements in queue
    Returns:
        image_batch: 4D tensor [batch_size, width, height, 3], dtype=tf.float32
        label_batch: 1D tensor [batch_size], dtype=tf.int32
    """

    image = tf.cast(image, tf.string)
    label = tf.cast(label, tf.int32)

    # make an input queue
    input_queue = tf.train.slice_input_producer([image, label])

    label = input_queue[1]
    image_contents = tf.read_file(input_queue[0])
    image = tf.image.decode_jpeg(image_contents, channels=3)

    image = tf.image.resize_image_with_crop_or_pad(image, image_W, image_H)

    # if you want to test the generated batches of images, you might want to comment the following line.
    image = tf.image.per_image_standardization(image)

    image_batch, label_batch = tf.train.batch([image, label],
                                              batch_size=batch_size,
                                              num_threads=64,
                                              capacity=capacity)
    # you can also use shuffle_batch
    #    image_batch, label_batch = tf.train.shuffle_batch([image,label],
    #                                                      batch_size=BATCH_SIZE,
    #                                                      num_threads=64,
    #                                                      capacity=CAPACITY,
    #                                                      min_after_dequeue=CAPACITY-1)

    label_batch = tf.reshape(label_batch, [batch_size])
    image_batch = tf.cast(image_batch, tf.float32)

    return image_batch, label_batch

    # TEST
    # To test the generated batches of images
    # When training the model, Do comment the following codes

    # import matplotlib.pyplot as plt
    #
    # BATCH_SIZE = 2
    # CAPACITY = 256
    # IMG_W = 208
    # IMG_H = 208
    #
    # train_dir = '/home/acrobat/DataSets/cats_vs_dogs/train/'
    # ratio = 0.2
    # tra_images, tra_labels, val_images, val_labels = get_files(train_dir, ratio)
    # tra_image_batch, tra_label_batch = get_batch(tra_images, tra_labels, IMG_W, IMG_H, BATCH_SIZE, CAPACITY)
    #
    #
    #
    # with tf.Session() as sess:
    #    i = 0
    #    coord = tf.train.Coordinator()
    #    threads = tf.train.start_queue_runners(coord=coord)
    #
    #    try:
    #        while not coord.should_stop() and i<1:
    #
    #            img, label = sess.run([tra_image_batch, tra_label_batch])
    #
    #            # just test one batch
    #            for j in np.arange(BATCH_SIZE):
    #                print('label: %d' %label[j])
    #                plt.imshow(img[j,:,:,:])
    #                plt.show()
    #            i+=1
    #
    #    except tf.errors.OutOfRangeError:
    #        print('done!')
    #    finally:
    #        coord.request_stop()
    #    coord.join(threads)


if __name__ == '__main__':
    tra_images, tra_labels, val_images, val_labels = get_files(train_dir, 0.2)
    train_batch, train_label_batch = get_batch(tra_images,
                                               tra_labels,
                                               208,
                                               208,
                                               64,
                                               2000)
